package com.thiswhyme.fsplus.sentinel.feign.instancefilter;

import com.google.common.base.Optional;
import com.netflix.loadbalancer.Server;
import com.netflix.loadbalancer.ZoneAvoidanceRule;
import org.springframework.util.CollectionUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.Set;

/**
 * @description: 自定义负载均衡规则
 * @author: thiswhyme
 * @date: 2021/4/29
 */
public class CustomZoneAvoidanceRule extends ZoneAvoidanceRule {

    /**
     * 保留无参构造解决ZoneAwareLoadBalancer.cloneRule时newInstance失败问题
     */
    public CustomZoneAvoidanceRule() {

    }

    public CustomZoneAvoidanceRule(BlockServerContext blockServerContext) {
        super();
        this.blockServerContext = blockServerContext;
    }

    private BlockServerContext blockServerContext;

    @Override
    public Server choose(Object key) {
        List<Server> servers = this.getLoadBalancer().getAllServers();

        //剔除断路的服务实例
        List<Server> filterServers = new ArrayList<>(servers.size());
        if(!CollectionUtils.isEmpty(servers)) {
            filterServers.addAll(servers);
            Set<String> unAvailableServers =  blockServerContext.getUnAvailableServers();
            if(!CollectionUtils.isEmpty(unAvailableServers)) {
                filterServers.removeIf(server -> unAvailableServers.contains(server.getHostPort()));
            }
        }
        //TODO 此处可根据可用实例域值添加告警逻辑

        Optional<Server> server = this.getPredicate().chooseRoundRobinAfterFiltering(filterServers, key);
        return server.isPresent() ? server.get() : null;
    }
}
